import sys
import os
import json
from typing import Optional, List
from dotenv import load_dotenv
from openai import OpenAI
import xml.etree.ElementTree as ET
from datetime import datetime
from colorama import init, Fore, Style
from halo import Halo
import logging
from logging.handlers import RotatingFileHandler

from tools import TOOLS, execute_tool_call, get_tool_call
from prompts import get_system_prompt

# 初始化日志记录器
def setup_logger():
    """设置日志记录器"""
    logger = logging.getLogger('ai_cli')
    logger.setLevel(logging.DEBUG)
    
    # 创建日志目录
    log_dir = 'logs'
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    
    # 文件处理器 - 按日期滚动
    log_file = os.path.join(log_dir, 'ai_cli.log')
    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=1024*1024,  # 1MB
        backupCount=5,
        encoding='utf-8'
    )
    file_handler.setLevel(logging.DEBUG)
    
    # 控制台处理器 - 仅在DEBUG模式下显示
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.DEBUG if os.getenv("DEBUG") else logging.ERROR)
    
    # 设置日志格式
    formatter = logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s',
        datefmt='%Y-%m-%d %H:%M:%S'
    )
    file_handler.setFormatter(formatter)
    console_handler.setFormatter(formatter)
    
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)
    
    return logger

# 初始化日志记录器
logger = setup_logger()

# 初始化colorama
init()

# 加载环境变量
load_dotenv()

# 检查必要的环境变量
required_env_vars = ["OPENAI_API_KEY", "OPENAI_BASE_URL", "MODEL"]
missing_vars = [var for var in required_env_vars if not os.getenv(var)]
if missing_vars:
    logger.error(f"缺少必要的环境变量: {', '.join(missing_vars)}")
    raise ValueError(f"缺少必要的环境变量: {', '.join(missing_vars)}")

# 配置OpenAI客户端
client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    base_url=os.getenv("OPENAI_BASE_URL")
)
MODEL = os.getenv("MODEL")
logger.info(f"使用模型: {MODEL}")

class ChatHistory:
    def __init__(self, file_path="chat_history.json"):
        self.file_path = file_path
        self.current_session = []
        self.load_history()

    def load_history(self):
        """加载历史记录"""
        if os.path.exists(self.file_path):
            try:
                with open(self.file_path, "r", encoding="utf-8") as f:
                    self.history = json.load(f)
            except:
                self.history = []
        else:
            self.history = []

    def save_history(self):
        """保存历史记录"""
        if self.current_session:
            session_data = {
                "timestamp": datetime.now().isoformat(),
                "messages": self.current_session
            }
            self.history.append(session_data)
            
            with open(self.file_path, "w", encoding="utf-8") as f:
                json.dump(self.history, f, ensure_ascii=False, indent=2)

    def add_message(self, role: str, content: str):
        """添加新消息"""
        self.current_session.append({
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        })

def print_colored(text: str, role: str = "system", end: str = "\n"):
    """打印彩色文本"""
    color = {
        "system": Fore.WHITE,
        "user": Fore.GREEN,
        "assistant": Fore.CYAN,
        "tool": Fore.YELLOW,
        "error": Fore.RED,
        "thinking": Fore.MAGENTA
    }.get(role, Fore.WHITE)
    
    print(f"{color}{text}{Style.RESET_ALL}", end=end, flush=True)

def process_command(command: str, history: ChatHistory) -> None:
    """处理单个命令"""
    spinner = Halo(text='思考中...', spinner='dots')
    try:
        # 初始化消息列表
        system_prompt = get_system_prompt()
        logger.debug("系统提示词已加载")
        
        messages = [{"role": "system", "content": system_prompt}]
        messages.append({"role": "user", "content": command})
        history.add_message("user", command)
        logger.info(f"用户输入: {command}")
        
        # 获取AI响应
        spinner.start()
        logger.debug("开始获取AI响应")
        response = client.chat.completions.create(
            model=MODEL,
            messages=messages,
            temperature=0,
            stream=True
        )
        spinner.stop()
        
        # 处理流式响应
        print_colored("\nAI: ", "assistant", end="")
        full_response = ""
        for chunk in response:
            if chunk.choices[0].delta.content:
                content = chunk.choices[0].delta.content
                print_colored(content, "assistant", end="")
                full_response += content
        print()  # 添加最后的换行
        
        logger.debug(f"AI完整响应: {full_response}")
        history.add_message("assistant", full_response)
        
        # 检查工具调用
        tool_call = get_tool_call(full_response)
        if tool_call:
            tool_name, xml_str = tool_call
            logger.info(f"检测到工具调用: {tool_name}")
            logger.debug(f"工具调用XML: {xml_str}")
            
            # 执行工具调用
            spinner.text = '执行中...'
            spinner.start()
            result = execute_tool_call(xml_str)
            spinner.stop()
            
            # 打印工具执行结果
            print_colored(f"\n工具执行结果: ", "tool", end="")
            print_colored(result, "tool")
            logger.info(f"工具执行结果: {result}")
            
            history.add_message("system", f"工具执行结果：\n{result}")
            messages.append({"role": "system", "content": f"工具执行结果：\n{result}"})
            
            # 将结果反馈给AI进行判断
            logger.debug("将结果反馈给AI进行判断")
            spinner.text = '分析结果中...'
            spinner.start()
            response = client.chat.completions.create(
                model=MODEL,
                messages=messages,
                temperature=0,
                stream=True
            )
            spinner.stop()
            
            # 处理AI的反馈
            print_colored("\nAI: ", "assistant", end="")
            feedback_response = ""
            for chunk in response:
                if chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    print_colored(content, "assistant", end="")
                    feedback_response += content
            print()
            
            logger.debug(f"AI反馈: {feedback_response}")
            history.add_message("assistant", feedback_response)
            
            # 如果是搜索结果，显示原始链接
            if tool_name == 'search_web':
                try:
                    result_data = json.loads(result)
                    if result_data.get('success') and result_data.get('results'):
                        logger.debug("处理搜索结果链接")
                        print_colored("\n相关链接：", "system")
                        for i, item in enumerate(result_data['results'], 1):
                            print_colored(f"{i}. {item['title']}", "system")
                            print_colored(f"   {item['url']}", "system")
                except Exception as e:
                    logger.error(f"处理搜索结果时出错: {str(e)}")
            
        else:
            logger.debug("未检测到工具调用")
            
    except Exception as e:
        if spinner.spinner_id:
            spinner.fail("AI响应出错")
        error_msg = f"错误：{str(e)}"
        logger.error(f"处理命令时出错: {str(e)}", exc_info=True)
        print_colored(error_msg, "error")
        history.add_message("error", error_msg)

def main(args: List[str]) -> None:
    """主函数"""
    history = ChatHistory()
    logger.info("AI CLI工具启动")
    
    try:
        # 显示欢迎信息
        print_colored("\nAI工具助手已启动（输入'退出'结束对话）\n", "system")
        logger.debug("显示欢迎信息")
        
        # 如果有命令行参数，先处理它
        if len(args) > 1:
            command = " ".join(args[1:])
            logger.info(f"处理命令行参数: {command}")
            process_command(command, history)
            print()
        
        # 进入交互循环
        while True:
            try:
                command = input("\n请输入命令: ").strip()
                if not command:
                    continue
                    
                if command.lower() in ['exit', 'quit', '退出']:
                    logger.info("用户请求退出")
                    print_colored("\n感谢使用，再见！", "system")
                    break
                    
                process_command(command, history)
                print()
                
            except KeyboardInterrupt:
                logger.info("用户中断程序")
                print_colored("\n\n程序已终止", "system")
                break
            except Exception as e:
                logger.error(f"处理用户输入时出错: {str(e)}", exc_info=True)
                print_colored(f"\n发生错误：{str(e)}", "error")
                continue
    
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
        print_colored("\n\n程序已终止", "system")
    except Exception as e:
        logger.error(f"程序运行时出错: {str(e)}", exc_info=True)
        print_colored(f"\n程序出错：{str(e)}", "error")
    finally:
        history.save_history()
        logger.info("会话结束，历史记录已保存")
        print_colored("\n会话结束", "system")

if __name__ == "__main__":
    main(sys.argv) 